import pandas as pd
import re
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


class Chatbot(object):
    def __init__(self):
        self.clf = TfidfVectorizer(ngram_range=(1, 3))

    @staticmethod
    def get_robots_train_data():
        """
        获取训练集数据
        :return:
        """
        df = pd.read_csv('./data/nscb.csv')
        return df['Postinner link']

    def get_data(self):
        """
        在原数据的基础上，切分问题和答案
        :return:
        """
        train_data = self.get_robots_train_data()

        clist = []
        for data in train_data:
            cparis = re.findall(": (.*?)(?:$|\n)", data)
            for info in list(zip(cparis, cparis[1:])):
                clist.append(info)

        return clist

    def get_similar_answer(self, train_data, question):
        """
        获取跟训练集问题最接近的那个答案
        :param question:
        :return:
        """
        vec = self.clf.fit_transform(train_data['question'])
        my_question = self.clf.transform([question])

        # 匹配跟问题最近接的答案
        cs = cosine_similarity(my_question, vec)
        rs = pd.Series(cs[0]).sort_values(ascending=False)

        return train_data.iloc[rs.index[0]]['answer']

    def run(self, question):
        # 获取数据
        clist = self.get_data()

        convo_frame = pd.Series(dict(clist)).to_frame().reset_index()
        convo_frame.columns = ['question', 'answer']

        return self.get_similar_answer(train_data=convo_frame, question=question)


my_question = 'Say goodbye, Clevercake'
charbot = Chatbot()
answer = charbot.run(question=my_question)
print(answer)
